import json

from src.tools.dms_operator import dmsoperator
from src.tools.chain import gen_chain
from src.tools.extract_jsonstring import re_json_string
from src.build_aitemplate.prompt import Prompt
from src.build_aitemplate.to_document import to_document_bpr, to_document_sop
from src.build_aitemplate.db import faiss_embedding
from src.tools.information_exchange import pub_gen

class BuildBPRAiTemplate:

    def __init__(self):
        self.prompt = Prompt()

    def find_key_values(self, target_key,data):
        found_values = []
        if isinstance(data, dict):
            for key, value in data.items():
                if key == target_key and value:
                    found_values.append(value)
                found_values.extend(self.find_key_values(value, target_key))
        elif isinstance(data, list):
            for item in data:
                found_values.extend(self.find_key_values(item, target_key))
        return found_values

    def parse_jsonfile(self, file: dict):
        # 根据文件路径获得对应数据信息
        json_data = dmsoperator.get_jsonfile(file.get("realPath"))
        # 获得建模段落信息
        parainfo = []
        for para in json_data["paragraphInfos"]:
            for para_text in para["textInfos"]:
                if para_text["highlightColor"] != "0, 0, 0, 0":
                    para_dict = {
                        "fileId": file.get("docId"),
                        "filePath": json_data["filePath"],
                        "fileName": file.get("documentName"),
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "text": para["text"],
                        "tableCellInfos": None
                    }
                    parainfo.append(para_dict)
                    break
        # 获得建模单元格信息
        cellinfo = []
        for table in json_data["tableInfos"]:
            for cell in table["tableCellInfos"]:
                is_highlightcolor_cell = False
                cell_dict = {
                    "fileId": file.get("docId"),
                    "filePath": json_data["filePath"],
                    "fileName": file.get("documentName"),
                    "index": table["index"],
                    "sectionIndex": table["sectionIndex"],
                    "bodyChildIndex": table["bodyChildIndex"],
                    "text": None,
                    "tableCellInfos": [{
                        "rowIndex": cell["rowIndex"],
                        "columnIndex": cell["columnIndex"],
                        "paragraphInfos": []
                    }]
                }
                for para in cell["paragraphInfos"]:
                    para_dict = {
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "text": para["text"]
                    }
                    cell_dict["tableCellInfos"][0]["paragraphInfos"].append(para_dict)
                    for para_text in para["textInfos"]:
                        if para_text["highlightColor"] != "0, 0, 0, 0":
                            is_highlightcolor_cell = True
                            break
                if is_highlightcolor_cell:
                    cellinfo.append(cell_dict)
        return parainfo, cellinfo

    def build_bpr_aitemplate(self, format_file: dict, retr, task_id: str):
        template_format_res = []
        # 解析格式文件
        format_parainfo, format_cellinfo = self.parse_jsonfile(format_file)
        # for template_file in template_files:
        #     template_file_jsondata = dmsoperator.get_jsonfile(template_file.get("realPath"))
        #     docus = to_document_bpr(template_file_jsondata, template_file)
        #     faiss_embedding.add_documents(docus)
        # retr = faiss_embedding.as_retriever(search_type="mmr", search_kwargs={"k": 3, "fetch_k": 30})

        for para in format_parainfo:
            print(para["text"])
            # 检索来源信息
            template_jsontexts = retr.invoke(para["text"])
            template_jsontexts_content = []
            for jsontext in template_jsontexts:
                template_jsontexts_content.append(jsontext.page_content)
            print(template_jsontexts)
            # 对检索结果进行重排
            rerank_query = self.prompt.rerank_prompt(para["text"], json.dumps(template_jsontexts_content, ensure_ascii=False))
            print(rerank_query)
            template_index = gen_chain(rerank_query)
            print(template_index)
            try:
                template_index = json.loads(re_json_string(template_index))
            except:
                template_index = {"position": [0]}
            print(template_index)
            # 来源元数据信息
            template_jsontext_meta = []
            msg_res = ""
            for index_num in template_index["position"]:
                try:
                    template_jsontext_meta.append(template_jsontexts[index_num].metadata)
                    msg_res += template_jsontexts_content[index_num] + "\n"
                except:
                    break
            pub_gen(task_id, msg_res, format_file.get("docId"), format_file.get("taskDocId"))
            # text_key = self.find_key_values("text", template_jsontext_meta)
            # for msg in text_key:
            #     pub_gen(task_id, msg)
            template_format_res.append({"mark": para, "source": template_jsontext_meta})

        for cell in format_cellinfo:
            cell_text = ""
            for para in cell["tableCellInfos"][0]["paragraphInfos"]:
                cell_text = cell_text + para["text"] + "\n"
            cell_text = cell_text.rstrip("\n")

            print(cell_text)
            # 检索来源信息
            template_jsontexts = retr.invoke(cell_text)
            template_jsontexts_content = []
            for jsontext in template_jsontexts:
                template_jsontexts_content.append(jsontext.page_content)
            print(template_jsontexts)
            # 对检索结果进行重排
            rerank_query = self.prompt.rerank_prompt(cell_text, json.dumps(template_jsontexts_content, ensure_ascii=False))
            print(rerank_query)
            template_index = gen_chain(rerank_query)
            try:
                template_index = json.loads(re_json_string(template_index))
            except:
                template_index = {"position": [0]}
            print(template_index)
            # 来源元数据信息
            template_jsontext_meta = []
            msg_res = ""
            for index_num in template_index["position"]:
                try:
                    template_jsontext_meta.append(template_jsontexts[index_num].metadata)
                    msg_res += template_jsontexts_content[index_num] + "\n"
                except:
                    break
            pub_gen(task_id, msg_res, format_file.get("docId"), format_file.get("taskDocId"))
            # text_key = self.find_key_values("text", template_jsontext_meta)
            # for msg in text_key:
            #     pub_gen(task_id, msg)
            template_format_res.append({"mark": cell, "source": template_jsontext_meta})

        return template_format_res



class BuildSOPAiTemplate:
    def __init__(self):
        self.prompt = Prompt()

    def find_key_values(self, target_key,data):
        found_values = []
        if isinstance(data, dict):
            for key, value in data.items():
                if key == target_key:
                    found_values.append(value)
                found_values.extend(self.find_key_values(value, target_key))
        elif isinstance(data, list):
            for item in data:
                found_values.extend(self.find_key_values(item, target_key))
        return found_values

    def parse_jsonfile(self, format_file):
        json_data = dmsoperator.get_jsonfile(format_file.get("realPath"))
        # 找到建模段落
        parainfo = []
        for para in json_data["paragraphInfos"]:
            for para_text in para["textInfos"]:
                if para_text["highlightColor"] != "0, 0, 0, 0":
                    para_dict = {
                        "fileId": format_file.get("docId"),
                        "filePath": json_data["filePath"],
                        "fileName": format_file.get("documentName"),
                        "type": 4,
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "text": para["text"],
                        "tableCellInfos": None
                    }
                    parainfo.append(para_dict)
                    break
        # 找到建模表格
        tableinfo = []
        for table in json_data["tableInfos"]:
            is_highlightcolor_table = False
            table_dict = {
                "fileId": format_file.get("docId"),
                "filePath": json_data["filePath"],
                "fileName": format_file.get("documentName"),
                "type": 4,
                "index": table["index"],
                "sectionIndex": table["sectionIndex"],
                "bodyChildIndex": table["bodyChildIndex"],
                "text": None,
                "tableCellInfos": []
            }
            for cell in table["tableCellInfos"]:
                cell_dict = {
                    "rowIndex": cell["rowIndex"],
                    "columnIndex": cell["columnIndex"],
                    "paragraphInfos": []
                }
                for para in cell["paragraphInfos"]:
                    para_dict = {
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "text": para["text"]
                    }
                    cell_dict["paragraphInfos"].append(para_dict)
                    for para_text in para["textInfos"]:
                        if para_text["highlightColor"] != "0, 0, 0, 0":
                            is_highlightcolor_table = True
                            break
                table_dict["tableCellInfos"].append(cell_dict)
            if is_highlightcolor_table:
                tableinfo.append(table_dict)
        return parainfo, tableinfo


    def build_sop_aitemplate(self, format_file: dict, retr, task_id:str):
        template_format_res = []
        # 解析格式文件
        format_parainfo, format_tableinfo = self.parse_jsonfile(format_file)
        # for template_file in template_files:
        #     template_file_jsondata = dmsoperator.get_jsonfile(template_file.get("realPath"))
        #     docus = to_document_sop(template_file_jsondata, template_file)
        #     faiss_embedding.add_documents(docus)
        # retr = faiss_embedding.as_retriever(search_type="mmr", search_kwargs={"k": 3, "fetch_k": 30})

        for para in format_parainfo:
            print(para["text"])
            # 检索来源信息
            template_jsontexts = retr.invoke(para["text"])
            template_jsontexts_content = []
            for jsontext in template_jsontexts:
                template_jsontexts_content.append(jsontext.page_content)
            print(template_jsontexts)
            # 对检索结果进行重排
            rerank_query = self.prompt.rerank_prompt(para["text"], str(template_jsontexts_content))
            print(rerank_query)
            template_index = gen_chain(rerank_query)
            print(template_index)
            try:
                template_index = json.loads(re_json_string(template_index))
            except:
                template_index = {"position": [0]}
            print(template_index)
            # 来源元数据信息
            template_jsontext_meta = []
            msg_res = ""
            for index_num in template_index["position"]:
                try:
                    template_jsontext_meta.append(template_jsontexts[index_num].metadata)
                    msg_res += template_jsontexts_content[index_num] + "\n"
                except:
                    break
            pub_gen(task_id, msg_res, format_file.get("docId"), format_file.get("taskDocId"))
            # text_key = self.find_key_values("text",template_jsontext_meta)
            # for msg in text_key:
            #     pub_gen(task_id, msg)
            template_format_res.append({"mark":para, "source":template_jsontext_meta})

        for table in format_tableinfo:
            table_text = ""
            for cell in table["tableCellInfos"]:
                cell_text = ""
                for para in cell["paragraphInfos"]:
                    cell_text += para["text"]
                table_text += cell_text + "|"

            print(table_text)
            # 检索来源信息
            template_jsontexts = retr.invoke(table_text)
            template_jsontexts_content = []
            for jsontext in template_jsontexts:
                template_jsontexts_content.append(jsontext.page_content)
            print(template_jsontexts)
            # 对检索结果进行重排
            rerank_query = self.prompt.rerank_prompt(table_text, str(template_jsontexts_content))
            print(rerank_query)
            template_index = gen_chain(rerank_query)
            try:
                template_index = json.loads(re_json_string(template_index))
            except:
                template_index = {"position": [0]}
            print(template_index)
            # 来源元数据信息
            template_jsontext_meta = []
            msg_res = ""
            for index_num in template_index["position"]:
                try:
                    template_jsontext_meta.append(template_jsontexts[index_num].metadata)
                    msg_res += template_jsontexts_content[index_num] + "\n"
                except:
                    break
            pub_gen(task_id, msg_res, format_file.get("docId"), format_file.get("taskDocId"))
            # text_key = self.find_key_values("text", template_jsontext_meta)
            # for msg in text_key:
            #     pub_gen(task_id, msg)
            template_format_res.append({"mark": table, "source": template_jsontext_meta})

        return template_format_res

















